# -*- coding: utf-8 -*-
"""
省份地理边界工具模块

此模块提供用于提取和处理中国省份地理边界数据的函数和类。
主要功能包括：
- 从shapefile中提取省份边界数据
- 计算省份的经纬度范围
- 判断点是否在省份内
- 提取省份内的网格点
"""

import os
import geopandas as gpd
import pandas as pd
import numpy as np
from shapely.geometry import Point
from typing import List, Dict, Any, Optional, Tuple, Union


def read_province_shapefile(shp_path: str, encoding: str = 'utf-8') -> gpd.GeoDataFrame:
    """
    读取省份shapefile文件并确保使用地理坐标系统(WGS84)
    
    参数:
        shp_path: shapefile文件路径
        encoding: 文件编码，默认为utf-8
        
    返回:
        包含省份数据的GeoDataFrame对象
    """
    try:
        # 读取shapefile文件
        gdf = gpd.read_file(shp_path, encoding=encoding)
        print(f"成功读取shapefile文件，共有{len(gdf)}个省份")
        
        # 确保坐标系统是地理坐标系统(WGS84)
        if gdf.crs and not gdf.crs.is_geographic:
            print(f"原始坐标系统: {gdf.crs}")
            gdf = gdf.to_crs("EPSG:4326")  # 转换为WGS84地理坐标系
            print(f"已转换为WGS84坐标系统: {gdf.crs}")
            
        return gdf
    except Exception as e:
        print(f"读取shapefile文件时出错: {e}")
        raise


def filter_province(gdf: gpd.GeoDataFrame, target_provinces: List[str], 
                   name_field: str = 'NAME') -> gpd.GeoDataFrame:
    """
    从GeoDataFrame中筛选指定的省份
    
    参数:
        gdf: 包含所有省份数据的GeoDataFrame
        target_provinces: 目标省份名称列表
        name_field: 省份名称字段，默认为'NAME'
        
    返回:
        筛选后的GeoDataFrame
    """
    filtered_gdf = gdf[gdf[name_field].isin(target_provinces)]
    print(f"找到{len(filtered_gdf)}个目标省份")
    
    if len(filtered_gdf) < len(target_provinces):
        found_provinces = filtered_gdf[name_field].tolist()
        missing_provinces = [p for p in target_provinces if p not in found_provinces]
        print(f"未找到以下省份: {missing_provinces}")
        
    return filtered_gdf


def get_province_bounds(province_gdf: gpd.GeoDataFrame, province: str, 
                       name_field: str = 'NAME') -> Tuple[float, float, float, float]:
    """
    获取指定省份的边界范围
    
    参数:
        province_gdf: 包含省份边界数据的GeoDataFrame
        province: 省份名称
        name_field: 省份名称字段，默认为'NAME'
        
    返回:
        省份边界范围(min_lon, min_lat, max_lon, max_lat)
    """
    try:
        # 获取省份几何对象
        province_geom = province_gdf[province_gdf[name_field] == province]['geometry'].iloc[0]
        
        # 获取省份边界范围
        bounds = province_geom.bounds  # 返回(minx, miny, maxx, maxy)
        min_lon, min_lat, max_lon, max_lat = bounds
        
        print(f"{province}省份边界范围: 经度({min_lon:.4f}, {max_lon:.4f}), 纬度({min_lat:.4f}, {max_lat:.4f})")
        return min_lon, min_lat, max_lon, max_lat
    except Exception as e:
        print(f"获取{province}省份边界范围失败: {e}")
        raise


def is_point_in_province(province_gdf: gpd.GeoDataFrame, province: str, 
                        lon: float, lat: float, name_field: str = 'NAME') -> bool:
    """
    判断点是否在指定省份内
    
    参数:
        province_gdf: 包含省份边界数据的GeoDataFrame
        province: 省份名称
        lon: 经度
        lat: 纬度
        name_field: 省份名称字段，默认为'NAME'
        
    返回:
        点是否在省份内
    """
    try:
        # 获取省份几何对象
        province_geom = province_gdf[province_gdf[name_field] == province]['geometry'].iloc[0]
        
        # 创建点对象
        point = Point(lon, lat)
        
        # 判断点是否在省份内
        return province_geom.contains(point)
    except Exception as e:
        print(f"判断点是否在{province}省份内失败: {e}")
        return False


def extract_province_grid_points(province_gdf: gpd.GeoDataFrame, lons: np.ndarray, lats: np.ndarray, 
                               name_field: str = 'NAME', provinces: Optional[List[str]] = None) -> Dict[str, Dict]:
    """
    提取指定省份的观测点经纬度数据
    
    参数:
        province_gdf: 省份边界GeoDataFrame
        lons: 经度数组
        lats: 纬度数组
        name_field: 省份名称字段，默认为'NAME'
        provinces: 要提取的省份名称列表，如为None则提取所有省份
        
    返回:
        包含各省份观测点经纬度数据的字典
    """
    # 如果未指定省份，则使用所有省份
    if provinces is None:
        provinces = province_gdf[name_field].tolist()
    else:
        # 筛选指定的省份
        filtered_gdf = filter_province(province_gdf, provinces, name_field)
        provinces = filtered_gdf[name_field].tolist()
    
    if not provinces:
        print("没有有效的省份可以提取")
        return {}
    
    # 创建结果字典
    result = {}
    
    # 对每个省份提取观测点
    for province in provinces:
        # 获取省份几何对象
        province_geom = province_gdf[province_gdf[name_field] == province]['geometry'].iloc[0]
        
        # 获取省份边界范围
        bounds = province_geom.bounds  # 返回(minx, miny, maxx, maxy)
        min_lon, min_lat, max_lon, max_lat = bounds
        
        # 提取在省份范围内的经纬度点
        province_lons = lons[(lons >= min_lon) & (lons <= max_lon)]
        province_lats = lats[(lats >= min_lat) & (lats <= max_lat)]
        
        # 创建经纬度网格点
        lon_grid, lat_grid = np.meshgrid(province_lons, province_lats)
        
        # 将网格点展平为一维数组
        grid_points_flat = np.column_stack((lon_grid.flatten(), lat_grid.flatten()))
        
        # 筛选真正在省份多边形内的点
        in_province_points = []
        for point in grid_points_flat:
            if province_geom.contains(Point(point[0], point[1])):
                in_province_points.append(point)
        
        # 转换为numpy数组
        if in_province_points:
            grid_points = np.array(in_province_points)
        else:
            grid_points = np.empty((0, 2))
        
        # 保存到结果字典
        result[province] = {
            'grid_points': grid_points,
            'lon_count': len(province_lons),
            'lat_count': len(province_lats),
            'total_points': len(grid_points)
        }
        
        print(f"省份: {province}, 经度点数: {len(province_lons)}, 纬度点数: {len(province_lats)}, 有效点数: {len(grid_points)}")
    
    return result


def extract_province_bounds(filtered_gdf: gpd.GeoDataFrame, 
                           name_field: str = 'NAME', 
                           fullname_field: Optional[str] = None) -> pd.DataFrame:
    """
    提取省份的经纬度边界范围
    
    参数:
        filtered_gdf: 筛选后的省份GeoDataFrame
        name_field: 省份名称字段，默认为'NAME'
        fullname_field: 省份全称字段，默认为None
        
    返回:
        包含省份边界数据的DataFrame
    """
    results = []
    
    # 提取每个省份的边界范围
    for idx, row in filtered_gdf.iterrows():
        province_name = row[name_field]
        bounds = row['geometry'].bounds  # 返回(minx, miny, maxx, maxy)
        
        result_dict = {
            '省份': province_name,
            '最小经度': bounds[0],
            '最小纬度': bounds[1],
            '最大经度': bounds[2],
            '最大纬度': bounds[3],
            '经度范围': f"{bounds[0]:.6f} - {bounds[2]:.6f}",
            '纬度范围': f"{bounds[1]:.6f} - {bounds[3]:.6f}"
        }
        
        # 如果提供了全称字段，则添加全称
        if fullname_field and fullname_field in row:
            result_dict['全称'] = row[fullname_field]
        else:
            result_dict['全称'] = province_name
            
        results.append(result_dict)
    
    return pd.DataFrame(results)


def filter_provinces(gdf, target_provinces, name_field='NAME'):
    """
    从GeoDataFrame中筛选指定的省份（与filter_province函数相同，为兼容性提供）
    
    参数:
        gdf: 包含所有省份数据的GeoDataFrame
        target_provinces: 目标省份名称列表
        name_field: 省份名称字段，默认为'NAME'
        
    返回:
        筛选后的GeoDataFrame
    """
    return filter_province(gdf, target_provinces, name_field)